from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import StreamingResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
import httpx
import uvicorn
import os
from pathlib import Path
import mimetypes
from typing import Optional
import json
from pydantic import BaseModel

from smart_choice.datanodes import DataNodes
from smart_choice.decisiontree import DecisionTree

class TreeRequest(BaseModel):
    json_data: dict
    payoff_fn_code: str

def json_to_datanodes(json_data, payoff_fn):
    nodes = DataNodes()
    for node in json_data["nodes"]:
        name = node["name"]
        node_type = node["type"]
        branches = node.get("branches", [])

        if node_type == "decision":
            # 决策节点
            branches_tuple = [
                (b["label"], float(b["value"]), b["next"])
                for b in branches
            ]
            nodes.add_decision(name=name, branches=branches_tuple, maximize=True)

        elif node_type == "chance":
            # 概率节点
            branches_tuple = [
                (b["label"], float(b["probability"]), float(b["value"]), b["next"])
                for b in branches
            ]
            nodes.add_chance(name=name, branches=branches_tuple)

        elif node_type == "terminal":
            # 终点节点
            nodes.add_terminal(name=name, payoff_fn=payoff_fn)

        else:
            raise ValueError(f"Unsupported node type: {node_type}")

    return nodes


class SmartServer:
    def __init__(self, static_dir: Optional[str] = None):
        """
        Initialize the SmartServer.

        Args:
            static_dir (str, optional): Path to the directory containing static files.
                                      If not provided, uses 'src/frontend/dist' relative to the package.
        """
        # Initialize MIME types
        mimetypes.add_type('application/javascript', '.js')
        mimetypes.add_type('application/javascript', '.mjs')
        mimetypes.add_type('text/css', '.css')
        mimetypes.add_type('text/html', '.html')

        self.app = FastAPI(title="Smart Server")

        # Set default static directory relative to this file if not provided
        if static_dir is None:
            # 尝试多个可能的路径
            possible_paths = [
                os.path.join(os.path.dirname(__file__), "src", "frontend", "dist"),
                os.path.join(os.path.dirname(__file__), "static"),
                "static",  # 当前工作目录下的 static
            ]
            
            for path in possible_paths:
                if os.path.exists(path):
                    static_dir = path
                    break
            else:
                # 如果都找不到，使用第一个路径作为默认值
                static_dir = possible_paths[0]

        self.static_dir = Path(static_dir)
        self.http_client = httpx.AsyncClient(follow_redirects=True)

        # Configure CORS
        self.app.add_middleware(
            CORSMiddleware,
            allow_origins=["*"],
            allow_credentials=True,
            allow_methods=["*"],
            allow_headers=["*"],
        )

        # Register routes
        self._setup_routes()

    def _setup_routes(self):
        """Setup all route handlers"""

        @self.app.get("/api/health")
        async def health_check():
            """Health check endpoint"""
            return {"status": "healthy"}

        @self.app.post("/api/generate-tree")
        async def generate_tree(request: TreeRequest):
            """Generate decision tree from JSON data and payoff function"""
            try:
                # 动态创建 payoff 函数
                local_vars = {}
                exec(request.payoff_fn_code, globals(), local_vars)
                payoff_fn = local_vars.get('payoff_fn')
                
                if not payoff_fn:
                    raise HTTPException(status_code=400, detail="payoff_fn function not found in the provided code")
                
                # 调用 json_to_datanodes 转换为 datanode
                nodes = json_to_datanodes(request.json_data, payoff_fn)
                
                # 创建决策树进行评估
                tree = DecisionTree(nodes=nodes)
                tree.evaluate()
                tree.rollback()
                
                # 生成 dot 文件
                dot = tree.plot()
                dot_content = dot.source
                
                # 可选：保存 dot 文件到本地
                dot.save("tree_output.dot")
                
                return {
                    "success": True,
                    "dot_content": dot_content,
                    "message": "Decision tree generated successfully"
                }
                
            except Exception as e:
                raise HTTPException(status_code=500, detail=f"Error generating decision tree: {str(e)}")

        @self.app.get("/proxy/{path:path}")
        async def proxy_request(path: str):
            """Proxy endpoint that forwards requests to target URL"""
            try:
                base_url = "https://example.com"  # Configurable in the future
                url = f"{base_url}/{path}"
                response = await self.http_client.get(url)
                return StreamingResponse(
                    content=response.iter_bytes(),
                    status_code=response.status_code,
                    headers=dict(response.headers)
                )
            except Exception as e:
                raise HTTPException(status_code=500, detail=str(e))

        @self.app.get("/api/data")
        async def get_data():
            """Example data endpoint"""
            return {
                "data": [
                    {"id": 1, "name": "Item 1"},
                    {"id": 2, "name": "Item 2"},
                ]
            }

        # Setup static file handling
        if self.static_dir and self.static_dir.exists():
            async def custom_file_response(path: str):
                """Handle static file responses with proper MIME types"""
                file_path = self.static_dir / path
                if not file_path.exists() or not file_path.is_file():
                    file_path = self.static_dir / "index.html"
                    if not file_path.exists():
                        raise HTTPException(
                            status_code=404,
                            detail=f"File not found and no index.html in {self.static_dir}"
                        )

                content_type = mimetypes.guess_type(str(file_path))[0] or 'application/octet-stream'
                return FileResponse(
                    path=str(file_path),
                    media_type=content_type,
                    headers={"Content-Type": content_type}
                )

            @self.app.get("/{full_path:path}")
            async def serve_static(full_path: str):
                """Serve static files with proper MIME types"""
                if full_path.startswith("api/") or full_path.startswith("proxy/"):
                    raise HTTPException(status_code=404, detail="Not Found")
                return await custom_file_response(full_path)
        else:
            # If no static directory is available, show a helpful message
            @self.app.get("/")
            async def root():
                return {
                    "message": "Smart Server is running",
                    "static_dir": str(self.static_dir),
                    "static_dir_exists": self.static_dir.exists() if self.static_dir else False,
                    "api_endpoints": [
                        "/api/health",
                        "/api/data",
                        "/proxy/{path}"
                    ]
                }

    def run(self, host: str = "0.0.0.0", port: int = 8000, **kwargs):
        """
        Run the server.

        Args:
            host (str): Host to bind to. Defaults to "0.0.0.0".
            port (int): Port to bind to. Defaults to 8000.
            **kwargs: Additional arguments to pass to uvicorn.run
        """
        print(f"Starting Smart Server...")
        print(f"Static files directory: {self.static_dir}")
        print(f"Static files available: {self.static_dir.exists() if self.static_dir else False}")
        print(f"API endpoints available at: http://{host}:{port}/api/")

        if kwargs.get('reload', False):
            print("Running in reload mode - automatic reloading enabled")

        uvicorn.run(self.app, host=host, port=port, **kwargs)

    def get_app(self) -> FastAPI:
        """
        Get the FastAPI application instance.

        Returns:
            FastAPI: The FastAPI application instance
        """
        return self.app

    async def shutdown(self):
        """Cleanup resources"""
        await self.http_client.aclose()

# Example usage and entry point
def main():
    server = SmartServer()
    server.run(reload=True)

if __name__ == "__main__":
    main()
